package cn.itcast.tags.ml.classification

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.{OneHotEncoder, OneHotEncoderEstimator, StringIndexer, VectorAssembler}
import org.apache.spark.ml.regression.{LinearRegression, LinearRegressionModel, LinearRegressionTrainingSummary}
import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.storage.StorageLevel

/**
 * @author: xu
 * @desc: 基于泰坦尼克号数据集，使用逻辑回归构建分类模型，评估预测模型
 */
object TitanicLrClassification {
  def main(args: Array[String]): Unit = {
    // 构建SparkSession实例对象，通过建造者模式创建
    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.stripSuffix("$"))
      .master("local[3]")
      .config("spark.sql.shuffle.partitions", "3")
      .getOrCreate()
    import spark.implicits._

    // TODO: 1. 读取原始数据集：泰坦尼克号数据集csv文件
    val rawTitanicDF: DataFrame = spark.read
      .option("header", "true")
      .option("inferSchema", "true")
      .csv("datas/titanic/train.csv")
    /*
      root
       |-- PassengerId: integer (nullable = true)
       |-- Survived: integer (nullable = true)
       |-- Pclass: integer (nullable = true)
       |-- Name: string (nullable = true)
       |-- Sex: string (nullable = true)
       |-- Age: double (nullable = true)
       |-- SibSp: integer (nullable = true)
       |-- Parch: integer (nullable = true)
       |-- Ticket: string (nullable = true)
       |-- Fare: double (nullable = true)
       |-- Cabin: string (nullable = true)
       |-- Embarked: string (nullable = true)
     */
    // rawTitanicDF.printSchema()
    // rawTitanicDF.show(10, truncate = false)

    // TODO: 2. 数据处理：特征工程（提取、转换、选择）
    // 2.1 Age年龄字段有缺省值，填充为年龄字段平均值
    val avgAge: Double = rawTitanicDF
      .select($"Age")
      .filter($"Age".isNotNull)
      .select(round(avg($"Age"), 2).as("avgAge"))
      .first()
      .getAs[Double](0)
    println(s"avgAge = ${avgAge}")
    val ageTitanicDF: DataFrame = rawTitanicDF
      .select(
        $"Survived".as("label"),
        $"Pclass", $"Sex", $"SibSp", $"Parch", $"Fare", $"Age",
        // 如果年龄为null就使用上面计算出来的平均年龄
        when($"Age".isNotNull, $"Age").otherwise(avgAge).as("defaultAge")
      )

    // 2.2 对Sex字段特征转换，使用StringIndexer和OneHotEncoder
    val indexer: StringIndexer = new StringIndexer()
      .setInputCol("Sex")
      .setOutputCol("sexIndex")
    val indexerTitanicDF: DataFrame = indexer.fit(ageTitanicDF).transform(ageTitanicDF)
    // male -> [1.0, 0.0]  female -> [0.0, 1.0]
    val encoder: OneHotEncoder = new OneHotEncoder()
      .setInputCol("sexIndex")
      .setOutputCol("sexVector")
      .setDropLast(false)
    val sexTitanicDF: DataFrame = encoder.transform(indexerTitanicDF)

    // 2.3 将特征值组合，使用VectorAssembler
    val assembler: VectorAssembler = new VectorAssembler()
      .setInputCols(Array(
        "Pclass", "sexVector", "SibSp", "Parch", "Fare", "defaultAge"
      ))
      .setOutputCol("features")
    val titanicDF: DataFrame = assembler.transform(sexTitanicDF)
    // titanicDF.printSchema()
    // titanicDF.show(20,truncate = false)

    // 2.4 将特征数据集划分为训练集和测试集，并将训练缓存
    val Array(trainingDF, testingDF) = titanicDF.randomSplit(Array(0.8, 0.2))
    trainingDF.cache().count()

    // TODO: 3. 使用逻辑回归算法，设置参数训练模型
    val lr: LogisticRegression = new LogisticRegression()
      // 设置特征和标签列名称
      .setLabelCol("label")
      .setFeaturesCol("features")
      .setPredictionCol("prediction") // 使用模型预测时，预测值的列名称
      // 设置为二分类，默认使用auto即可
      .setFamily("binomial")
      // 是否对特征数据，进行标准化转换处理
      .setStandardization(true)
      // 设置算法相关的超参数值
      .setMaxIter(100)
      .setRegParam(0.1) // 正则化参数
      .setElasticNetParam(0.9) // 弹性网络参数
    val lrModel: LogisticRegressionModel = lr.fit(trainingDF)
    // 打印逻辑回归的斜率和截距，就是k和b
    println(s"Coefficients: ${lrModel.coefficients} Intercept: ${lrModel.intercept}")

    // TODO: 4. 模型评估与预测
    val predictionDF: DataFrame = lrModel.transform(testingDF)
    predictionDF
      .select($"label", $"prediction", $"probability", $"features")
      .show(30, truncate = false)
    val evaluator: MulticlassClassificationEvaluator = new MulticlassClassificationEvaluator()
      .setLabelCol("label")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")
    // accuracy = 0.8031914893617021
    println(s"accuracy = ${evaluator.evaluate(predictionDF)}")

    // 应用结束，关闭资源
    spark.stop()
  }
}